package com.cicc.wm.lock.jdbc;

import com.cicc.wm.lock.LockRepository;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanInitializationException;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.SmartInitializingSingleton;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.dao.DataIntegrityViolationException;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.support.DefaultTransactionDefinition;
import org.springframework.transaction.support.TransactionTemplate;
import org.springframework.util.Assert;

import javax.sql.DataSource;
import java.sql.Timestamp;
import java.time.Duration;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.UUID;

/**
 * jdbc锁操作类
 */
@Slf4j
public class JdbcLockRepository implements LockRepository, InitializingBean, DisposableBean,
        ApplicationContextAware, SmartInitializingSingleton {

    /**
     * 默认表前缀
     */
    public static final String DEFAULT_TABLE_PREFIX = "INT_";
    /**
     * 默认Region
     */
    public static final String DEFAULT_REGION = "DEFAULT";

    // SQL语句模板集中管理
    private static final String INSERT_QUERY_TEMPLATE = "INSERT INTO %sLOCK (REGION, LOCK_KEY, CLIENT_ID, CREATED_DATE, THREAD_ID, TIMEOUT, "
            + "EXPIRE_AT) VALUES (?, ?, ?, ?, ?, ?, ?)";
    private static final String IS_LOCKED_TEMPLATE = "SELECT COUNT(1) FROM %sLOCK "
            + "WHERE REGION=? AND LOCK_KEY=? AND CLIENT_ID=? AND THREAD_ID = ? AND EXPIRE_AT>=?";
    private static final String RENEW_LOCK_TEMPLATE = "UPDATE %sLOCK " +
            "SET EXPIRE_AT = DATE_ADD(NOW(6), INTERVAL (TIMEOUT * 1000) MICROSECOND) " +
            "WHERE REGION = ? AND LOCK_KEY = ? AND CLIENT_ID = ? AND THREAD_ID = ? AND EXPIRE_AT >= ?";

    private static final String REENTRY_LOCK_TEMPLATE = "UPDATE %sLOCK " +
            "SET EXPIRE_AT = DATE_ADD(NOW(6), INTERVAL (TIMEOUT * 1000) MICROSECOND), REENTRY_COUNT = REENTRY_COUNT + 1 " +
            "WHERE REGION = ? AND LOCK_KEY = ? AND CLIENT_ID = ? AND THREAD_ID = ? AND EXPIRE_AT >= ?";
    private static final String EXIT_LOCK_TEMPLATE = "UPDATE %sLOCK " +
            "SET REENTRY_COUNT = REENTRY_COUNT - 1 " +
            "WHERE REGION = ? AND LOCK_KEY = ? AND CLIENT_ID = ? AND THREAD_ID = ? AND EXPIRE_AT >= ? AND REENTRY_COUNT > 1";
    private static final String DELETE_QUERY_TEMPLATE = "DELETE FROM %sLOCK " +
            "WHERE REGION=? AND LOCK_KEY=? AND CLIENT_ID=? AND THREAD_ID=? AND REENTRY_COUNT = 1";
    private static final String DELETE_ALL_QUERY_TEMPLATE = "DELETE FROM %sLOCK WHERE REGION=? AND CLIENT_ID=?";
    private static final String PREEMPT_EXPIRED_LOCK_TEMPLATE = "UPDATE %sLOCK SET CLIENT_ID=?, CREATED_DATE=?, THREAD_ID = ?, TIMEOUT = ?, " +
            "EXPIRE_AT = ?, REENTRY_COUNT = 1 WHERE REGION=? AND LOCK_KEY=? AND EXPIRE_AT<?";

    private final String clientId;
    private final JdbcTemplate template;

    @Setter
    private String prefix = DEFAULT_TABLE_PREFIX;
    @Setter
    private String region = DEFAULT_REGION;

    // SQL语句实例
    private String preemptExpiredLockSql;
    private String renewLockSql;
    private String insertLockSql;
    private String isLockedSql;
    private String reentryLockSql;
    private String exitLockSql;
    private String unLockSql;
    private String unLockAllSql;
    private ApplicationContext applicationContext;
    private PlatformTransactionManager transactionManager;
    private TransactionTemplate defaultTransactionTemplate;
    private TransactionTemplate readOnlyTransactionTemplate;
    private TransactionTemplate readCommittedTransactionTemplate;

    public JdbcLockRepository(DataSource dataSource) {
        this(dataSource, UUID.randomUUID().toString());
    }

    public JdbcLockRepository(DataSource dataSource, String clientId) {
        Assert.hasText(clientId, "id must not be null nor empty");
        this.template = new JdbcTemplate(dataSource);
        this.clientId = clientId;
    }

    @Override
    public void afterPropertiesSet() {
        // 格式化SQL语句
        this.preemptExpiredLockSql = String.format(PREEMPT_EXPIRED_LOCK_TEMPLATE, this.prefix);
        this.reentryLockSql = String.format(REENTRY_LOCK_TEMPLATE, this.prefix);
        this.insertLockSql = String.format(INSERT_QUERY_TEMPLATE, this.prefix);

        this.isLockedSql = String.format(IS_LOCKED_TEMPLATE, this.prefix);
        this.renewLockSql = String.format(RENEW_LOCK_TEMPLATE, this.prefix);
        this.exitLockSql = String.format(EXIT_LOCK_TEMPLATE, this.prefix);
        this.unLockSql = String.format(DELETE_QUERY_TEMPLATE, this.prefix);
        this.unLockAllSql = String.format(DELETE_ALL_QUERY_TEMPLATE, this.prefix);
    }

    @Override
    public void afterSingletonsInstantiated() {
        if (this.transactionManager == null) {
            try {
                this.transactionManager = this.applicationContext.getBean(PlatformTransactionManager.class);
            } catch (BeansException ex) {
                throw new BeanInitializationException("A unique or primary 'PlatformTransactionManager' bean "
                        + "must be present in the application context.", ex);
            }
        }
        DefaultTransactionDefinition transactionDefinition =
                new DefaultTransactionDefinition(TransactionDefinition.PROPAGATION_REQUIRES_NEW);
        this.defaultTransactionTemplate = new TransactionTemplate(this.transactionManager, transactionDefinition);
        // It is safe to reuse the transactionDefinition - the TransactionTemplate makes copy of its properties.
        transactionDefinition.setReadOnly(true);
        this.readOnlyTransactionTemplate = new TransactionTemplate(this.transactionManager, transactionDefinition);
        transactionDefinition.setReadOnly(false);
        transactionDefinition.setIsolationLevel(TransactionDefinition.ISOLATION_READ_COMMITTED);
        this.readCommittedTransactionTemplate = new TransactionTemplate(this.transactionManager, transactionDefinition);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    /**
     * 关闭当前实例，删除所有锁
     */
    @Override
    public void close() {
        this.defaultTransactionTemplate.executeWithoutResult(
                transactionStatus -> {
                    int deleted = this.template.update(this.unLockAllSql, this.region, this.clientId);
                    log.info("[JdbcLockRepository] 关闭实例，已删除锁数量: {}，region={}, clientId={}", deleted, this.region, this.clientId);
                });
    }


    /**
     * 尝试获取锁
     *
     * @param threadId  线程ID
     * @param lockKey   锁Key
     * @param leaseTime 持有时长（毫秒）
     * @return 是否获取成功
     */
    @Override
    public boolean acquire(long threadId, String lockKey, long leaseTime) {
        Assert.hasText(lockKey, "lockKey must not be empty");
        Assert.isTrue(leaseTime > 0, "leaseTime must be positive");
        Boolean result = this.readCommittedTransactionTemplate.execute(transactionStatus -> {
            // 执行加锁操作
            try {
                // 锁重入
                int reentryCount = this.template.update(this.reentryLockSql, this.region, lockKey,
                        this.clientId, threadId, epochMillis());
                if (reentryCount > 0) {
                    log.debug("[JdbcLockRepository] 锁重入，lockKey={}, threadId={}", lockKey, threadId);
                    return true;
                }

                // 锁过期,直接抢占
                int preemptCount = this.template.update(this.preemptExpiredLockSql, this.clientId, epochMillis(), threadId, leaseTime,
                        this.ttlEpochMillis(Duration.ofMillis(leaseTime)), this.region, lockKey, epochMillis());
                if (preemptCount > 0) {
                    log.debug("[JdbcLockRepository] 锁过期,抢占锁成功，lockKey={}, threadId={}", lockKey, threadId);
                    return true;
                }

                // 创建锁
                boolean inserted = this.template.update(this.insertLockSql, this.region, lockKey, this.clientId, epochMillis(),
                        threadId, leaseTime, this.ttlEpochMillis(Duration.ofMillis(leaseTime))) > 0;
                if (inserted) {
                    log.debug("[JdbcLockRepository] 加锁成功，lockKey={}, threadId={}", lockKey, threadId);
                }
                return inserted;
            } catch (DataIntegrityViolationException ex) {
                log.debug("[JdbcLockRepository] 加锁失败，lockKey={}, threadId={}, 原因: {}", lockKey, threadId, ex.getMessage());
                return false;
            }
        });
        return Boolean.TRUE.equals(result);
    }

    /**
     * 判断锁是否被当前线程持有
     *
     * @param threadId 线程ID
     * @param lockKey  锁Key
     * @return 是否持有
     */
    @Override
    public boolean isAcquired(long threadId, String lockKey) {
        Assert.hasText(lockKey, "lockKey must not be empty");
        final Boolean result = this.readOnlyTransactionTemplate.execute(
                transactionStatus -> Integer.valueOf(1).equals(this.template.queryForObject(this.isLockedSql,
                        Integer.class, this.region, lockKey, this.clientId, threadId, this.epochMillis())));
        if (Boolean.TRUE.equals(result)) {
            log.debug("[JdbcLockRepository] 锁已被当前线程持有，lockKey={}, threadId={}", lockKey, threadId);
        } else {
            log.debug("[JdbcLockRepository] 锁未被当前线程持有，lockKey={}, threadId={}", lockKey, threadId);
        }
        return Boolean.TRUE.equals(result);
    }

    /**
     * 续约锁
     *
     * @param threadId 线程ID
     * @param lockKey  锁Key
     * @return 是否续约成功
     */
    @Override
    public boolean renew(long threadId, String lockKey) {
        Assert.hasText(lockKey, "lockKey must not be empty");
        final Boolean result = this.defaultTransactionTemplate.execute(transactionStatus -> {
            int rowsAffected = this.template.update(this.renewLockSql, this.region, lockKey, this.clientId,
                    threadId, this.epochMillis());
            if (rowsAffected == 0) {
                log.warn("[JdbcLockRepository] 续约失败，lockKey={}, threadId={}", lockKey, threadId);
            } else {
                log.info("[JdbcLockRepository] 续约成功，lockKey={}, threadId={}", lockKey, threadId);
            }
            return rowsAffected != 0;
        });
        return Boolean.TRUE.equals(result);
    }

    /**
     * 删除当前线程锁信息
     *
     * @param threadId 线程ID
     * @param lockKey  锁Key
     */
    @Override
    public boolean delete(long threadId, String lockKey) {
        Assert.hasText(lockKey, "lockKey must not be empty");
        return Boolean.TRUE.equals(this.defaultTransactionTemplate.execute(
                transactionStatus -> {
                    // 重入锁，退出时减少重入次数
                    int reentryCount = this.template.update(this.exitLockSql, this.region, lockKey,
                            this.clientId, threadId, this.epochMillis());
                    if (reentryCount > 0) {
                        log.debug("[JdbcLockRepository] 退出锁，lockKey={}, threadId={}", lockKey, threadId);
                        return true;
                    } else {
                        int deleted = this.template.update(this.unLockSql, this.region, lockKey, this.clientId, threadId);
                        log.info("[JdbcLockRepository] 删除锁，lockKey={}, threadId={}, deleted={}", lockKey, threadId, deleted);
                        return deleted > 0;
                    }
                }));
    }

    /**
     * 销毁当前实例，删除所有锁
     */
    @Override
    public void destroy() throws Exception {
        this.defaultTransactionTemplate.executeWithoutResult(
                transactionStatus -> {
                    int deleted = this.template.update(this.unLockAllSql, this.region, this.clientId);
                    log.info("[JdbcLockRepository] destroy()，已删除锁数量: {}，region={}, clientId={}", deleted, this.region, this.clientId);
                });
    }

    /**
     * 计算锁的过期时间戳
     *
     * @param duration 持有时长
     * @return 过期时间戳
     */
    private Timestamp ttlEpochMillis(Duration duration) {
        return Timestamp.valueOf(currentTime().plus(duration));
    }

    /**
     * 获取当前时间戳
     *
     * @return 当前时间戳
     */
    private Timestamp epochMillis() {
        return Timestamp.valueOf(currentTime());
    }

    /**
     * 获取当前UTC时间
     *
     * @return 当前UTC时间
     */
    private LocalDateTime currentTime() {
        return LocalDateTime.now(ZoneOffset.UTC);
    }
}
